from typing import List, Dict, Optional
import numpy as np
import os
import networkx as nx
from sklearn.metrics import mutual_info_score
import traceback

from app.config.config import settings
from app.tools.redis_client import redis_client
from app.schemas.facenet_base_schema import AlgorithmResponse, OutputParams, TaskStatus, AlgorithmMiddleResponse
from app.schemas.facenet_base_schema import FacetNetInputParams, FacetNetOutputParams, TimeStepResult
from app.utils.logger import logger
from app.utils.llm_client import LLMClient

class FacetNetService:
    def __init__(self):
        self.temp_dir = os.path.join(settings.TEMP_DIR, "facetnet")
        os.makedirs(self.temp_dir, exist_ok=True)
    
    async def update_task_status(self, task_id: str, task_data: dict, progress: int, logs: str):
        """更新任务状态"""
        try:
            response = AlgorithmMiddleResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.RUNNING,
                task_progress=progress,
                task_logs=logs,
                input_params=task_data.get("input_params"),
                metrics=[]
            )
            # 使用 model_dump() 确保数据可序列化
            await redis_client.update_data(f'{task_id}_result', response.model_dump())
            logger.info(f"任务状态更新成功: {task_id}, 进度: {progress}%, 日志: {logs}")
        except Exception as e:
            logger.error(f"任务状态更新失败: {task_id}, 错误: {str(e)}")
            raise
    
    def Sim(self, adj_mat, weighted=False):
        """计算相似度矩阵"""
        if weighted:
            return adj_mat
        else:
            N = adj_mat.shape[0]
            W = np.zeros((N, N))
            for i in range(N):
                for j in range(N):
                    if adj_mat[i, j] > 0:
                        W[i, j] = 1
            return W
    
    def soft_modularity(self, soft_comm, W):
        """计算软模块度"""
        N, K = soft_comm.shape
        Q = 0
        W_sum = np.sum(W)
        for i in range(N):
            for j in range(N):
                for k in range(K):
                    Q += (W[i, j] - np.sum(W[i, :]) * np.sum(W[:, j]) / W_sum) * soft_comm[i, k] * soft_comm[j, k]
        return Q / W_sum
    
    def param_update(self, X, A, Y, W, alpha):
        """参数更新"""
        W_apprx = X * A * X.T
        N, M = Y.shape
        X_new, A_new = np.zeros(X.shape), np.zeros(A.shape)
        for k in range(M):
            for i in range(N):
                for j in range(N):
                    X_new[i, k] += W[i, j] * A[k, k] * X[j, k] / W_apprx[i, j] if W_apprx[i, j] != 0 else 0
                    A_new[k, k] += W[i, j] * X[i, k] * X[j, k] / W_apprx[i, j] if W_apprx[i, j] != 0 else 0
                X_new[i, k] *= (2 * alpha * X[i, k])
                A_new[k, k] *= (alpha * A[k, k])
                X_new[i, k] += (1 - alpha) * Y[i, k]
                A_new[k, k] += (1 - alpha) * Y[i, k]
        X_new = np.matrix(X_new / np.sum(X_new, axis=0).reshape(1, M)) if np.sum(X_new) > 0 else X_new
        A_new = np.matrix(A / np.sum(A_new)) if np.sum(A_new) > 0 else A_new
        Y = X_new * A_new
        return X_new, A_new, Y
    
    def read_edgelist(self, edges, weighted=False):
        """读取边列表并构建邻接矩阵"""
        idmap = set()
        edge_cache = {}
        
        for edge in edges:
            if weighted:
                u, v, w = edge
            else:
                u, v = edge[:2]
                w = 1.0
            
            edge_cache[(u, v)] = w
            idmap.add(u)
            idmap.add(v)
        
        idmap = list(idmap)                                   # 数组下标与结点唯一id标识的映射
        idmap_inv = {nid: i for i, nid in enumerate(idmap)}   # 结点唯一id标识与数组下标的映射
        N = len(idmap)
        adj_mat = np.zeros((N, N))
        
        for (u, v), w in edge_cache.items():
            adj_mat[idmap_inv[u], idmap_inv[v]] = w
        
        adj_mat += adj_mat.T
        return idmap, idmap_inv, adj_mat
    
    def facetnet_algorithm(self, edges_list, alpha, tsteps, N, M, with_truth=False, truth_comms=None):
        """FacetNet算法实现"""
        results = []
        community_assignments = []  # 存储每个时间步的节点社区分配
        
        X, A = np.random.rand(N, M), np.diag(np.random.rand(M))
        X, A = np.matrix(X / np.sum(X, axis=0).reshape(1, M)), np.matrix(A / np.sum(A))
        Y = X * A
        
        for t in range(tsteps):
            # 读取当前时间步的边列表
            idmap, idmap_inv, adj_mat = self.read_edgelist(edges_list[t], weighted=False)
            
            # 如果提供了真实社区信息
            comm_map = {}
            if with_truth and truth_comms and t < len(truth_comms):
                for node_id, comm_id in truth_comms[t]:
                    comm_map[node_id] = comm_id
            
            W = self.Sim(adj_mat, weighted=False)
            X_new, A_new, Y = self.param_update(X, A, Y, W, alpha)
            
            D = np.zeros((N,))
            for i in range(N):
                D[i] = np.sum(Y[i, :])
            
            D = np.matrix(np.diag(D))
            soft_comm = D.I * X_new * A_new
            
            # 获取节点社区分配
            comm_pred = np.array(np.argmax(soft_comm, axis=1)).ravel()
            
            # 将当前时间步的社区分配添加到结果列表
            community_assignments.append(comm_pred.tolist())
            
            # 重新计算社区关系矩阵，使其更清晰反映社区结构
            community_matrix = np.zeros((N, N))
            for i in range(N):
                for j in range(N):
                    # 如果两个节点属于同一社区，则设置为1
                    if comm_pred[i] == comm_pred[j]:
                        community_matrix[i][j] = 1
                    else:
                        # 不同社区的节点之间，计算他们所属社区的关联强度
                        # 可以设置为0或一个较小的值，这里我们设置为0
                        community_matrix[i][j] = 0
            
            time_result = TimeStepResult(
                time=t,
                communities=comm_pred.tolist(),
                soft_modularity=float(self.soft_modularity(soft_comm, W)),
                community_matrix=community_matrix.tolist()  # 保留社区关系矩阵
            )
            
            if with_truth and truth_comms and t < len(truth_comms):
                comm = np.array([comm_map.get(idmap[i], 0) for i in range(N)])
                time_result.mutual_info = float(mutual_info_score(comm, comm_pred))
            
            results.append(time_result)
            X, A = X_new, A_new
        
        return results, community_assignments
    
    async def run_facetnet(self, input_params: dict, task_id: str) -> AlgorithmResponse:
        """运行FacetNet算法的主函数"""
        task_data = await redis_client.get_data(task_id)
        input_params = task_data.get("input_params")
        
        try:
            # 从输入参数中提取FacetNet参数
            facetnet_params = FacetNetInputParams(**input_params.get("facetnet_params", {}))
            
            await self.update_task_status(task_id, task_data, 10, "开始运行FacetNet算法...")
            
            # 运行算法
            np.random.seed(0)  # 设置随机种子以确保结果可重现
            results, community_assignments = self.facetnet_algorithm(
                facetnet_params.edges_list,
                facetnet_params.alpha,
                facetnet_params.tsteps,
                facetnet_params.N,
                facetnet_params.M,
                facetnet_params.with_truth,
                facetnet_params.truth_comms
            )
            
            await self.update_task_status(task_id, task_data, 90, f"FacetNet算法运行完成，共处理了{facetnet_params.tsteps}个时间步")
            
            # 构建输出参数
            facetnet_output = FacetNetOutputParams(
                results=results,
                algorithm="FacetNet",
                parameters={
                    "alpha": facetnet_params.alpha,
                    "tsteps": facetnet_params.tsteps,
                    "N": facetnet_params.N,
                    "M": facetnet_params.M
                },
                community_results=community_assignments  # 使用节点社区分配结果
            )
            
            # 返回最终结果
            return AlgorithmResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=OutputParams(facetnet_results=facetnet_output),
                metrics=[]
            )
        except Exception as e:
            error_msg = f"错误: {str(e)}\n{traceback.format_exc()}"
            logger.error(error_msg)
            
            # 更新任务状态为失败
            error_response = AlgorithmMiddleResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.FAILED,
                task_progress=0,
                task_logs=error_msg,
                input_params=task_data.get("input_params"),
                metrics=[]
            )
            await redis_client.update_data(f'{task_id}_result', error_response.model_dump())
            raise 